import os
from datasets import load_dataset

def dump_images_and_labels(data,split):
    # 从数据集中获取指定的数据集划分
    data = data[split]
    for i,example in enumerate(data):
        # 获取图像和标签数据
        image = example["image"]
        labels = example["litter"]["label"]
        bboxes = example["litter"]["bbox"]
        targets = []
        for label,box in zip(labels,bboxes):
            # 将标签和边界框信息转换为目标字符串
            targets.append(f"{label} {box[0]} {box[1]} {box[2]} {box[3]}")
        with open(f"datasets/labels/{split}/{i}.txt","w") as f:
            # 将目标字符串写入标签文件
            for target in targets:
                f.write(target + "\n")
                
        # 保存图像文件
        image.save(f"datasets/images/{split}/{i}.png")

if __name__ == "__main__":
    # 加载数据集
    dataset = load_dataset("kili-technology/plastic_in_river")
    # 创建目录用于存储图像和标签文件
    os.makedirs("datasets/images/train", exist_ok=True)
    os.makedirs("datasets/images/validation", exist_ok=True)
    os.makedirs("datasets/labels/train", exist_ok=True)
    os.makedirs("datasets/labels/validation", exist_ok=True)
    
    # 转换并保存训练集和验证集的图像和标签
    dump_images_and_labels(dataset,"train")
    dump_images_and_labels(dataset,"validation")